import { CCBoolean, CCFloat, CCString, Component, _decorator } from "cc";
import { DemonstrationWriter } from "./DemonstrationWriter";
import { Agent } from "../Agent";
import { BehaviorParameters } from "../Policies/BehaviorParameters";
const { ccclass, menu, property } = _decorator;


/**
 * 
 * @en
 *  The Demonstration Recorder component facilitates the recording of demonstrations
 * used for imitation learning.
 * 
 * Add this component to the [GameObject] containing an {@link Agent}
 * to enable recording the agent for imitation learning. You must implement the
 * {@link Agent.heuristic} function of the agent to provide manual control
 * in order to record demonstrations. 
 * See [Imitation Learning - Recording Demonstrations] for more information. 
 * 
 * Note: This component will only record demonstrations when running in the Cocos Editor.
 *
 * @zh
 * 演示记录器组件有助于记录用于模仿学习的演示。
 * 
 * 
 * 将此组件添加到包含 {@link Agent} 的 [GameObject] 中，以启用为模仿学习记录代理。
 * 您必须实现代理的 {@link Agent.heuristic} 函数，以提供手动控制以记录演示。
 * 有关更多信息，请参见[模仿学习 - 记录演示]。
 * 
 * 注意：只有在 COCOS 编辑器中运行时，此组件才会记录演示。
 * 
 */
@ccclass("DemonstrationRecorder")
@menu("ML Agents/DemonstrationRecorder")
export class DemonstrationRecorder extends Component {

    /**
     * @en
     * Whether or not to record demonstrations.
     * @zh
     * 是否记录演示。
     */
    @property(CCBoolean)
    private _record: boolean = false;
    public get record(): boolean {
        return this._record;
    }
    public set record(value: boolean) {
        this._record = value;
    }

    /**
     * @en
     * Number of steps to record. The editor will stop playing when it reaches this threshold.
     * Set to zero to record indefinitely.
     * @zh
     * 要记录的步数。当达到此阈值时，编辑器将停止播放。
     */
    @property(CCFloat)
    private _numStepsToRecord: number = 0;
    public get numStepsToRecord(): number {
        return this._numStepsToRecord;
    }
    public set numStepsToRecord(value: number) {
        this._numStepsToRecord = value;
    }

    /**
     * @en
     * Base demonstration file name. If multiple files are saved, the additional filenames
     * will have a sequence of unique numbers appended.
     * 
     * @zh
     * 基本演示文件名。如果保存多个文件，将附加一系列唯一数字的附加文件名。
     * 
     */
    @property(CCString)
    private _demonstrationName: string = '';
    public get demonstrationName(): string {
        return this._demonstrationName;
    }
    public set demonstrationName(value: string) {
        this._demonstrationName = value;
    }

    /**
     * @en
     *  Directory to save the demo files. Will default to a "Demonstrations/" folder in the
     * Application data path if not specified.
     * 
     * @zh
     * 保存演示文件的目录。如果未指定，将默认为项目数据路径中的“Demonstrations/”文件夹。
     */
    @property(CCString)
    private _demonstrationDirectory: string = '';
    public get demonstrationDirectory(): string {
        return this._demonstrationDirectory;
    }
    public set demonstrationDirectory(value: string) {
        this._demonstrationDirectory = value;
    }

    private _demoWriter: DemonstrationWriter | null = null!;
    private static readonly MaxNameLength: number = 16;
    private static readonly k_ExtensionType: string = ".demo";
    private static readonly k_DefaultDirectoryName: string = "Demonstrations";
    private _agent: Agent | null = null!;


    onEnable(): void {
        this._agent = this.getComponent(Agent);
    }

    update(): void {
        if (!this._record) {
            return;
        }

        this.lazyInitialize();
        // if (this._numStepsToRecord > 0 && this._demoWriter.numSteps >= this._numStepsToRecord) {

        // }
    }

    private lazyInitialize(): DemonstrationWriter {
        if (this._demoWriter != null) {
            return this._demoWriter;
        }

        if (this._agent == null) {
            this._agent = this.getComponent(Agent);
        }

        const behaviorParams = this.getComponent(BehaviorParameters);
        if (this._demonstrationName == null || this._demonstrationName === "") {
            this._demonstrationName = behaviorParams?.behaviorName ?? "Demonstration";
        }
        if (this._demonstrationDirectory == null || this._demonstrationDirectory === "") {
            this._demonstrationDirectory = DemonstrationRecorder.k_DefaultDirectoryName;
        }

        this._demonstrationName = DemonstrationRecorder.sanitizeName(this._demonstrationName, DemonstrationRecorder.MaxNameLength);
        this._demoWriter = new DemonstrationWriter(this._demonstrationName);
        this.addDemonstrationWriterToAgent(this._demoWriter);
        return this._demoWriter;
    }

    /**
     * @en
     * Removes all characters except alphanumerics from demonstration name.
     * Shorten name if it is longer than the maxNameLength.
     * 
     * @zh
     * 从演示名称中删除除字母数字之外的所有字符。
     * 如果名称超过 maxNameLength，则缩短名称。
     * 
     * 
     * @param demoName 
     * @param maxNameLength 
     * @returns 
     */
    private static sanitizeName(demoName: string, maxNameLength: number): string {
        const rgx = new RegExp("[^a-zA-Z0-9 -]");
        demoName = demoName.replace(rgx, "");
        if (demoName.length > maxNameLength) {
            demoName = demoName.substring(0, maxNameLength);
        }
        return demoName;
    }


    public close(): void {
        if (this._demoWriter != null) {
            this.removeDemonstrationWriterFromAgent(this._demoWriter);
            this._demoWriter.close();
            this._demoWriter = null;
        }
    }

    onDestroy(): void {
        this.close();
    }

    public addDemonstrationWriterToAgent(demoWriter: DemonstrationWriter): void {
        const behaviorParams = this.getComponent(BehaviorParameters);
        if (behaviorParams) {
            demoWriter.initialize(
                this._demonstrationName,
                behaviorParams?.brainParameters,
                behaviorParams?.fullyQualifiedBehaviorName
            );
        }

        this._agent?.demonstrationWriters.add(demoWriter);
    }

    public removeDemonstrationWriterFromAgent(demoWriter: DemonstrationWriter): void {
        this._agent?.demonstrationWriters.delete(demoWriter);
    }
}